import warnings
from collections import OrderedDict
from .logger import get_logger, coloring
from torchmetrics import Metric
from typing import Union
import torch
from torch import Tensor

logger = get_logger("paddlevideo")

__all__ = ['AverageMeter', 'build_record', 'log_batch', 'log_epoch']


class AverageMeter(object):
    """
    Computes and stores the average and current value
    """

    def __init__(self, name='', fmt='f', need_avg=True):
        self.name = name
        self.fmt = fmt
        self.need_avg = need_avg
        self.reset()

    def reset(self):
        """ reset """
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        """ update """
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count

    @property
    def total(self):
        return '{self.name}_sum: {self.sum:{self.fmt}}'.format(self=self)

    @property
    def total_minute(self):
        return '{self.name}_sum: {s:{self.fmt}} min'.format(s=self.sum / 60,
                                                            self=self)

    @property
    def mean(self):
        return '{self.name}_avg: {self.avg:{self.fmt}}'.format(
            self=self) if self.need_avg else ''

    @property
    def value(self):
        return '{self.name}: {self.val:{self.fmt}}'.format(self=self)


def build_record(cfg, ddp_used=False):
    if ddp_used:  # pytorch ddp
        record_list = [("loss", AverageMeter('loss', '7.4f'))]  # add loss
    else:  # no ddp
        record_list = [("batch_time", AverageMeter('batch_cost', '.4f')),
                       ("lr", AverageMeter('lr', 'f')),
                       ("loss", AverageMeter('loss', '7.4f'))]  # add loss
    # add model metric
    if 'Recognizer1D' in cfg.framework:  # TODO: required specify str in framework
        record_list.append(("hit_at_one", AverageMeter("hit_at_one", '.4f')))
        record_list.append(("perr", AverageMeter("perr", '.4f')))
        record_list.append(("gap", AverageMeter("gap", '.4f')))
    elif 'Recognizer' in cfg.framework:
        record_list.append(("Acc@1", AverageMeter("Acc@1", '.4f')))
        record_list.append(("Acc@1", AverageMeter("Acc@5", '.4f')))

    record_list = OrderedDict(record_list)
    return record_list


def log_batch(metric_list: OrderedDict, batch_id: int, epoch_id: int, total_epoch: int,
              mode: str, ips: str):
    batch_cost = str(metric_list['batch_time'].value) + ' sec,'
    metric_values = []
    for m in metric_list:
        if m != 'batch_time':
            metric_values.append(metric_list[m].value)

    metric_str = ' '.join([str(v) for v in metric_values])
    num_digits = len(str(total_epoch // 1))
    # fmt第一次调用format时以{{指定输出含{ , 如num_digits=2，第一步输出{:>2d}/{:<2d}
    # 第二次调用format则利用前一次指定输出的{做为替换字段标识，第二步输出 epoch_id/total_epoch
    epoch_str = '{{:>{0}d}}/{{:<{0}d}}'.format(num_digits).format(epoch_id, total_epoch)
    epoch_str = "[{:s}]: epoch:[{}]".format(mode.title(), epoch_str)
    step_str = " step:{:<4d}".format(batch_id)

    logger.info("{:s} {:s} {:s} {:s} {}".format(
        coloring(epoch_str, "HEADER") if batch_id == 0 else epoch_str,
        coloring(step_str, "PURPLE"), coloring(metric_str, 'OKGREEN'),
        coloring(batch_cost, "OKGREEN"), ips))


def log_epoch(metric_list: OrderedDict, epoch: int, mode: str, ips: str):
    batch_cost = 'avg_' + str(metric_list['batch_time'].mean) + ' sec,'
    epoch_sum = str(metric_list['batch_time'].total) + ' sec,'

    metric_values = []
    for m in metric_list:
        if not (m == 'batch_time' or m == 'lr'):
            metric_values.append(metric_list[m].mean)
    metric_str = ' '.join([str(v) for v in metric_values])

    end_epoch_str = "[{:s}]  END epoch:{:<3d}".format(mode.title(), epoch)

    logger.info("{:s} {:s} {:s} {:s} {}".format(
        coloring(end_epoch_str, "RED"), coloring(metric_str, "OKGREEN"),
        coloring(batch_cost, "OKGREEN"), coloring(epoch_sum, "OKGREEN"), ips))


def log_ddp_batch(metric_list: OrderedDict, batch_id: int,
                  epoch_id: int, total_epoch: int, mode: str):
    metric_values = []
    for key, value in metric_list.items():
        if key != 'batch_time':
            metric_values.append('{}:{:.4f}'.format(key, value))
    metric_str = ' '.join([str(v) for v in metric_values])

    num_digits = len(str(total_epoch // 1))
    # fmt第一次调用format时以{{指定输出含{ , 如num_digits=2，第一步输出{:>2d}/{:<2d}
    # 第二次调用format则利用前一次指定输出的{做为替换字段标识，第二步输出 epoch_id/total_epoch
    epoch_str = '{{:>{0}d}}/{{:<{0}d}}'.format(num_digits).format(epoch_id, total_epoch)
    epoch_str = "[{:s}]: epoch:[{}]".format(mode.title(), epoch_str)
    step_str = " step:{:<4d}".format(batch_id)

    logger.info("{:s} {:s} {:s}".format(
        coloring(epoch_str, "HEADER") if batch_id == 0 else epoch_str,
        coloring(step_str, "PURPLE"), coloring(metric_str, 'OKGREEN')))


def log_ddp_epoch(metric_list: OrderedDict, epoch: int, mode: str):
    metric_values = []
    for key, value in metric_list.items():
        if key != 'batch_time':
            metric_values.append('{}_avg:{:.4f}'.format(key, value))
    metric_str = ' '.join([str(v) for v in metric_values])
    end_epoch_str = "[{:s}]  END epoch:{:<3d}".format(mode.title(), epoch)
    logger.info("{:s} {:s}".format(coloring(end_epoch_str, "RED"),
                                   coloring(metric_str, "OKGREEN")))
